#include "arpcache.h"
#include "arp.h"
#include "ether.h"
#include "icmp.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <signal.h>
#include <assert.h>
#include <log.h>

#include<ip.h>

static arpcache_t arpcache;

// initialize IP->mac mapping, request list, lock and sweep thread
void arpcache_init()
{
	bzero(&arpcache, sizeof(arpcache_t));

	init_list_head(&(arpcache.req_list));

	pthread_mutex_init(&arpcache.lock, NULL);

	pthread_create(&arpcache.thread, NULL, arpcache_sweep, NULL);
}

// release all the resources when exiting
void arpcache_destroy()
{
	pthread_mutex_lock(&arpcache.lock);

	struct arp_req *req_entry = NULL, *req_q;
	list_for_each_entry_safe(req_entry, req_q, &(arpcache.req_list), list) {
		struct cached_pkt *pkt_entry = NULL, *pkt_q;
		list_for_each_entry_safe(pkt_entry, pkt_q, &(req_entry->cached_packets), list) {
			list_delete_entry(&(pkt_entry->list));
			free(pkt_entry->packet);
			free(pkt_entry);
		}

		list_delete_entry(&(req_entry->list));
		free(req_entry);
	}

	pthread_kill(arpcache.thread, SIGTERM);

	pthread_mutex_unlock(&arpcache.lock);
}

// look up the IP->mac mapping, need pthread_mutex_lock/unlock
// Traverse the table to find whether there is an entry with the same IP and mac address with the given arguments.
int arpcache_lookup(u32 ip4, u8 mac[ETH_ALEN])
{
	// assert(0 && "TODO: function arpcache_lookup not implemented!");
	pthread_mutex_lock(&arpcache.lock);

	struct arp_cache_entry *entry = NULL;
	// 遍历缓存表，找到一个过期的entry
	int i;
	for (i = 0; i < MAX_ARP_SIZE; i++) {
		if (arpcache.entries[i].valid && arpcache.entries[i].ip4==ip4) {
			// 解锁 放入mac
			memcpy(mac,arpcache.entries[i].mac,ETH_ALEN);
			pthread_mutex_unlock(&arpcache.lock);
			return 1;
		}
	}
	// 没找到，返回0
	pthread_mutex_unlock(&arpcache.lock);
	return 0;
}

// insert the IP->mac mapping into arpcache, need pthread_mutex_lock/unlock
// If there is a timeout entry (attribute valid in struct) in arpcache, replace it.
// If there isn't a timeout entry in arpcache, randomly replace one.
// If there are pending packets waiting for this mapping, fill the ethernet header for each of them, and send them out.
// Tips:
// arpcache_t是完整的arp缓存表，里边的req_list是一个链表，它的每个节点(用arp_req结构体封装)里又存着一个链表头，这些二级链表(节点类型是cached_pkt)缓存着相同目标ip但不知道mac地址的包
void arpcache_insert(u32 ip4, u8 mac[ETH_ALEN])
{
    log(DEBUG, "lock the arpcache");
    pthread_mutex_lock(&arpcache.lock);

    struct arp_cache_entry *entry = NULL;

    // 遍历缓存表，找到一个过期的 entry
    int i;
    for (i = 0; i < MAX_ARP_SIZE; i++) {
        if (arpcache.entries[i].valid == 0) {
            entry = &arpcache.entries[i];
            break;
        }
    }

    // 如果没有过期的 entry，随机替换一个
    if (entry == NULL) {
        log(DEBUG, "No expired entry found, replacing a random entry");
        srand(time(NULL));
        i = rand() % MAX_ARP_SIZE;
        entry = &arpcache.entries[i];
    } else {
        log(DEBUG, "Find expired entry, update it");
    }

    // 更新 entry
    entry->ip4 = ip4;
    assert(mac != NULL);
    memcpy(entry->mac, mac, ETH_ALEN);
    entry->added = time(NULL);
    entry->valid = 1;

    log(DEBUG, "Insert IP->MAC done: " IP_FMT " -> " ETHER_STRING,
        HOST_IP_FMT_STR(entry->ip4), ETHER_FMT(entry->mac));

    // 遍历 req_list，找到所有等待这个 IP 的包
    struct arp_req *req_entry = NULL, *req_q = NULL;
    list_for_each_entry_safe(req_entry, req_q, &(arpcache.req_list), list) {
        if (req_entry->ip4 == ip4) {
            // 保存所有等待此 IP 的包信息
            log(DEBUG, "Saving all packets waiting for this IP " IP_FMT,
                HOST_IP_FMT_STR(ip4));

            // 计算需要处理的包数量
            int packet_count = 0;
            struct cached_pkt *pkt_entry = NULL, *pkt_q = NULL;
            list_for_each_entry_safe(pkt_entry, pkt_q, &(req_entry->cached_packets), list) {
                packet_count++;
            }

            // 分配临时数组
            char **packets = malloc(packet_count * sizeof(char *));
            int *packet_lens = malloc(packet_count * sizeof(int));
            int index = 0;

            list_for_each_entry_safe(pkt_entry, pkt_q, &(req_entry->cached_packets), list) {
                // 为每个包分配新内存并复制内容
                packets[index] = malloc(pkt_entry->len);
                if (packets[index] == NULL) {
                    log(ERROR, "Memory allocation failed for packets[%d]", index);
                    exit(1); // 或者更优雅地退出
                }

				if(!pkt_entry&&!pkt_entry->packet){
					log(ERROR,"not found the packet");
				}

				
                memcpy(packets[index], pkt_entry->packet, pkt_entry->len);
                packet_lens[index] = pkt_entry->len;
                index++;

                // 删除链表节点并释放原始内存
                list_delete_entry(&(pkt_entry->list));
                free(pkt_entry->packet);
                free(pkt_entry);
            }

            // 删除 ARP 请求项
            list_delete_entry(&(req_entry->list));
            free(req_entry);

            // 解锁
            pthread_mutex_unlock(&arpcache.lock);

            // 发送所有 ICMP 包
            for (int i = 0; i < packet_count; i++) {
                log(DEBUG, "Sending the %d-th ICMP packet", i);

                // 封装icmp_packet
                struct ether_header* eh=(struct ether_header*)(packets[i]);
                
                memcpy(eh->ether_dhost,mac,ETH_ALEN);

                struct iphdr* ip=packet_to_ip_hdr(packets[i]);
                struct icmphdr *icmp = (struct icmphdr *)IP_DATA(ip);

                ip_send_packet(packets[i],packet_lens[i]);
                // free(packets[i]); // 释放每个包的内存
            }

            // 释放临时数组
            free(packets);
            free(packet_lens);

            // 重新加锁继续处理其他项
            pthread_mutex_lock(&arpcache.lock);
            continue;
        }
    }

    log(DEBUG, "unlock arpcache");
    pthread_mutex_unlock(&arpcache.lock);
}


// append the packet to arpcache
// Look up in the list which stores pending packets, if there is already an entry with the same IP address and iface, 
// which means the corresponding arp request has been sent out, just append this packet at the tail of that entry (The entry may contain more than one packet).
// Otherwise, malloc a new entry with the given IP address and iface, append the packet, and send arp request.
// Tips:
// arpcache_t是完整的arp缓存表，里边的req_list是一个链表，它的每个节点(类型是arp_req)里又存着一个链表头，这些二级链表(节点类型是cached_pkt)缓存着相同目标ip但不知道mac地址的包
void arpcache_append_packet(iface_info_t *iface, u32 ip4, char *packet, int len)
{
    // 分配 cached_pkt 的内存，同时包含包数据的空间
    pthread_mutex_lock(&arpcache.lock);
    log(DEBUG,"lock the arpcache. start append packet");
    struct cached_pkt *cached_packet = malloc(sizeof(struct cached_pkt));
    if (!cached_packet) {
        log(ERROR, "Memory allocation failed for cached_pkt");
        pthread_mutex_unlock(&arpcache.lock);
        return; // 或者处理内存分配失败的情况
    }

    // 填充结构体的字段
	
	// 分配 packet 数据的内存
    cached_packet->packet = malloc(len);
    memcpy(cached_packet->packet, packet, len); // 将包数据拷贝到结构体中的 packet 字段
    cached_packet->len = len;

    struct iphdr *ip=packet_to_ip_hdr(cached_packet->packet);
    // 准备更改icmp_hdr的一些字段，其他都复制
    struct icmphdr *icmp_hdr = (struct icmphdr *)IP_DATA(ip);
    


    struct arp_req *req_entry, *tmp;
    list_for_each_entry_safe(req_entry, tmp, &(arpcache.req_list), list) {
        // 如果找到了相同的 IP 地址和 iface
        if (req_entry->iface == iface && req_entry->ip4 == ip4) {
            // 对应的 ARP 请求已经发出，只需要将包加入到该 entry 的尾部
            log(DEBUG,"ARP request has been sent, directly add the packet");
            list_add_tail(&cached_packet->list, &req_entry->cached_packets);
            pthread_mutex_unlock(&arpcache.lock);
            return;
        }
    }

    // 如果没有找到，创建一个新的 ARP 请求 entry
    log(DEBUG,"Create new request entry");
    struct arp_req *new_entry = malloc(sizeof(struct arp_req));
    if (!new_entry) {
        log(ERROR, "Memory allocation failed for arp_req");
        free(cached_packet); // 清理之前分配的内存
        pthread_mutex_unlock(&arpcache.lock);
        return;
    }

    new_entry->iface = iface;
    new_entry->ip4 = ip4;
    init_list_head(&new_entry->cached_packets);
    list_add_tail(&cached_packet->list, &new_entry->cached_packets);

    // 将新的请求加入 ARP 请求队列
    list_add_tail(&new_entry->list, &arpcache.req_list);
    log(DEBUG,"newly append packet, sending request");

    new_entry->retries++;
    new_entry->sent = time(NULL);

    arp_send_request(new_entry->iface,new_entry->ip4);
   
    pthread_mutex_unlock(&arpcache.lock);
    return ;
}


// sweep arpcache periodically
// for IP->mac entry, if the entry has been in the table for more than 15 seconds, remove it from the table
// for pending packets, if the arp request is sent out 1 second ago, while the reply has not been received, retransmit the arp request
// If the arp request has been sent 5 times without receiving arp reply, for each pending packet, send icmp packet (DEST_HOST_UNREACHABLE), and drop these packets
// tips
// arpcache_t是完整的arp缓存表，里边的req_list是一个链表，它的每个节点(类型是arp_req)里又存着一个链表头，这些二级链表(节点类型是cached_pkt)缓存着相同目标ip但不知道mac地址的包
void *arpcache_sweep(void *arg) 
{
	while (1) {
		sleep(1);
		// assert(0 && "TODO: function arpcache_sweep not implemented!");
		pthread_mutex_lock(&arpcache.lock);
		log(DEBUG,"sweep lock");

		// 遍历路由表arpcache_t，先查找IP->mac映射表
		struct arp_cache_entry *entry = NULL;
		for (int i = 0; i < MAX_ARP_SIZE; i++) {
			entry = &arpcache.entries[i];
			// 如果entry过期，删除
			if (entry->valid && time(NULL) - entry->added > ARP_ENTRY_TIMEOUT) {
				entry->valid = 0;
				// 这样删除就好了，以后insert会直接取代
			}

		}
		
		// 遍历req_list，查找pending packets
		struct arp_req *req_entry = NULL, *req_q;
		list_for_each_entry_safe(req_entry, req_q, &(arpcache.req_list), list) {
			// 如果req_entry过期，删除
            log(DEBUG,"start to sweep the pending packets");
            log(INFO,"this packet sent by %d, now the time is %d, that is %d ago",req_entry->sent,time(NULL),time(NULL)-req_entry->sent);
			if (time(NULL) - req_entry->sent > 1) {
				// 如果重试次数超过5次，发送icmp包
                log(INFO,"this req entry has been retransmit: %d times",req_entry->retries);
				if (req_entry->retries >= ARP_REQUEST_MAX_RETRIES) {
					struct cached_pkt *pkt_entry = NULL, *pkt_q;
					list_for_each_entry_safe(pkt_entry, pkt_q, &(req_entry->cached_packets), list) {
						// 对所有pending packets发送icmp error包，即目的主机不可达

						// 防止递归上锁死锁，先解锁
						pthread_mutex_unlock(&arpcache.lock);

						log(DEBUG,"unreachable! sending icmp packet");
						icmp_send_packet(pkt_entry->packet, pkt_entry->len, ICMP_DEST_UNREACH, ICMP_HOST_UNREACH);

						// 继续上锁
						pthread_mutex_lock(&arpcache.lock);

						list_delete_entry(&(pkt_entry->list));
					}
					list_delete_entry(&(req_entry->list));
				} else {
					// 否则重发arp请求
					req_entry->retries++;
					req_entry->sent = time(NULL);
					arp_send_request(req_entry->iface, req_entry->ip4);
				}
			}
		}
        log(DEBUG,"arpcache sweep: unlock");
		pthread_mutex_unlock(&arpcache.lock);
	}
	pthread_mutex_unlock(&arpcache.lock);
	return NULL;
}
